import os.path
import tkinter as tk
from tkinter import filedialog
import re
import json

#变量
coorXYZ = []
Allatom = []
#读取原子 坐标
def ReadAtomAndCoor(text):
    for line in text:
        atom = line.split(" ")[0]
        Allatom.append(atom)
        OXYZ = re.findall(r"\-?\d+\.?\d*",line)
        #XYZ=[float(OXYZ[1]),float(OXYZ[2]),float(OXYZ[3])]
        XYZ=[("%.8f" % float(OXYZ[1])),("%.8f" % float(OXYZ[2])),("%.8f" % float(OXYZ[3]))]
        coorXYZ.append(XYZ)
        print(Allatom)
    return Allatom
#写入优化的gjf文件
def WriteOPTInput(text):
    ReadAtomAndCoor(text)
    with open(file_folder+r'OPT.gjf', 'w') as output:
        output.write(r'%chk='+file_folder+'OPT.chk'+'\n')
        output.write(r'%nprocshared='+Core+'\n')
        output.write(r'%mem='+Mem+'MW'+'\n')
        output.write(Basis+Basis2+Basis3+Basis4+'\n')
        output.write('\n')
        output.write(r'[No Title]'+'\n')
        output.write('\n')
        output.write(Charge+r' '+Spin)
        output.write('\n')
        for x in range(len(Allatom)):
            output.write(r" ")
            output.write(Allatom[x])
            if len(Allatom[x]) == 1:
                output.write(r' ')
            if float(coorXYZ[x][0]) < 0:
                output.write(r'                 '+str(coorXYZ[x][0]))
            else:
                output.write(r'                  '+str(coorXYZ[x][0]))
            if float(coorXYZ[x][1]) < 0:
                output.write(r'   '+str(coorXYZ[x][1]))
            else:
                output.write(r'    '+str(coorXYZ[x][1]))
            if float(coorXYZ[x][2]) < 0:
                output.write(r'   '+str(coorXYZ[x][2])+'\n')
            else:
                output.write(r'    '+str(coorXYZ[x][2])+'\n')

#写入能量的gjf文件
def ReadAtomAndCoor_Energy(text):
    LineNum1 = 0
    LastLineNum1 = 0
    chosentext=[]
    #找出最后一行 Standard orientation: 的行数
    for line in text:
        search = re.findall(r'Standard orientation:',line)
        if search != []:
            LastLineNum1 = LineNum1
        LineNum1 = LineNum1 + 1

    LineNum2 = 0
    LastLineNum2 = 0
    #找出最后一行 Population analysis using the SCF density. 的行数
    for line in text:
        search = re.findall(r'Population analysis using the SCF density.',line)
        if search != []:
            LastLineNum2 = LineNum2
        LineNum2 = LineNum2 + 1

    chosentext = text[LastLineNum1:LastLineNum2]

    #分割出数字坐标段落
    count3 = 0
    chosentext2 = []
    for line in chosentext:
        search = re.findall(r'---------------------------------------------------------------------', line)
        if search != []:
            count3 += 1
        elif count3 == 2:
            chosentext2.append(line)
    for line in chosentext2:
        OXYZ = re.findall(r"\-?\d+\.?\d*", line)
        XYZ = [("%.8f" % float(OXYZ[3])),("%.8f" % float(OXYZ[4])),("%.8f" % float(OXYZ[5]))]
        coorXYZ.append(XYZ)

    #提取原子
    for line in lines2:
        atom = line.split(" ")[0]
        Allatom.append(atom)
    #写入文件
    with open(file_folder+r'ENERGY.gjf', 'w') as output:
        output.write(r'%chk='+file_folder+'ENERGY.chk'+'\n')
        output.write(r'%nprocshared='+Core+'\n')
        output.write(r'%mem='+Mem+'MW'+'\n')
        output.write(Basis+Basis2+Basis3+Basis4+'\n')
        output.write('\n')
        output.write(r'[No Title]'+'\n')
        output.write('\n')
        output.write(r'0 1')
        output.write('\n')
        for x in range(len(Allatom)):
            output.write(r" ")
            output.write(Allatom[x])
            if len(Allatom[x]) == 1:
                output.write(r' ')
            if float(coorXYZ[x][0]) < 0:
                output.write(r'                '+str(coorXYZ[x][0]))
            else:
                output.write(r'                 '+str(coorXYZ[x][0]))
            if float(coorXYZ[x][1]) < 0:
                output.write(r'   '+str(coorXYZ[x][1]))
            else:
                output.write(r'    '+str(coorXYZ[x][1]))
            if float(coorXYZ[x][2]) < 0:
                output.write(r'   '+str(coorXYZ[x][2])+'\n')
            else:
                output.write(r'    '+str(coorXYZ[x][2])+'\n')

if __name__ == '__main__':
    #打开文件
    root = tk.Tk()
    root.withdraw()
    file_path = filedialog.askopenfilename(filetypes=[('gjf输入文件 out优化输出文', '.gjf .out')])
    #决定任务、电荷
    Charge = input('输入电荷数 如为默认值 直接按enter')
    if Charge:
        int(Charge)
    if os.path.splitext(file_path)[-1] == ".gjf":
        namep = -len(file_path.split(r"/")[-1])
        file_folder = file_path[:namep]
        with open(file_path, 'r') as f:
            lines = f.readlines()[5:]
            lines.pop()
        #读取高斯计算参数
        p = open('parameter1.json', 'r')
        content = p.read()
        Gau_Parameter = json.loads(content)
        Basis = '# opt b3lyp/'
        Basis2 = Gau_Parameter['Basis2']
        Basis3 = Gau_Parameter['Basis3']
        Basis4 = Gau_Parameter['Basis4']
        Core = Gau_Parameter['Core']
        Mem = Gau_Parameter['Mem']
        if Charge == False:
            Charge = Gau_Parameter['Charge']
        Spin = Gau_Parameter['Spin']
        #写入优化文件
        WriteOPTInput(lines)

    if os.path.splitext(file_path)[-1] == ".out":
        namep = -len(file_path.split(r"/")[-1])
        file_folder = file_path[:namep]
        opt_file_path = file_folder + 'Untitled-1.gjf'
        print(opt_file_path)
        with open(file_path, 'r') as f:
            lines = f.readlines()[5:]
        with open(opt_file_path, 'r') as f1:
            lines2 = f1.readlines()[5:]
            lines2.pop()
        p = open('parameter2.json', 'r')
        content = p.read()
        Gau_Parameter = json.loads(content)
        Basis = '# b3lyp/'
        Basis2 = Gau_Parameter['Basis2']
        Basis3 = Gau_Parameter['Basis3']
        Basis4 = Gau_Parameter['Basis4']
        Core = Gau_Parameter['Core']
        Mem = Gau_Parameter['Mem']
        #写入能量文件
        ReadAtomAndCoor_Energy(lines)